{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# SageMaker V3 Distributed Local Training Example\n",
    "\n",
    "This notebook demonstrates how to run distributed training locally using SageMaker V3 ModelTrainer with multiple Docker containers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import subprocess\n",
    "import tempfile\n",
    "import shutil\n",
    "import numpy as np\n",
    "\n",
    "from sagemaker.train.model_trainer import ModelTrainer, Mode\n",
    "from sagemaker.train.configs import SourceCode, Compute, InputData\n",
    "from sagemaker.train.distributed import Torchrun\n",
    "from sagemaker.core.helper.session_helper import Session"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NOTE: Local mode requires Docker to be installed and running.\n",
    "import os\n",
    "os.environ['PATH'] = '/usr/local/bin:/Applications/Docker.app/Contents/Resources/bin:' + os.environ['PATH']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 1: Setup Session and Create Test Data\n",
    "\n",
    "Initialize the SageMaker session and create the necessary test data and training script."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sagemaker_session = Session()\n",
    "DEFAULT_CPU_IMAGE = \"763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training:2.0.0-cpu-py310\"\n",
    "\n",
    "# Create temporary directories\n",
    "temp_dir = tempfile.mkdtemp()\n",
    "source_dir = os.path.join(temp_dir, \"source\")\n",
    "data_dir = os.path.join(temp_dir, \"data\")\n",
    "train_dir = os.path.join(data_dir, \"train\")\n",
    "test_dir = os.path.join(data_dir, \"test\")\n",
    "\n",
    "os.makedirs(source_dir, exist_ok=True)\n",
    "os.makedirs(train_dir, exist_ok=True)\n",
    "os.makedirs(test_dir, exist_ok=True)\n",
    "\n",
    "print(f\"Created temporary directories in: {temp_dir}\")\n",
    "print(\"Note: This will use multiple Docker containers locally for distributed training!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 2: Create Training Data and Scripts\n",
    "\n",
    "Generate the test data and training scripts needed for distributed local training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create test data\n",
    "np.random.seed(42)\n",
    "x_train = np.random.randn(100, 4).astype(np.float32)\n",
    "y_train = np.random.randn(100).astype(np.float32)\n",
    "x_test = np.random.randn(20, 4).astype(np.float32)\n",
    "y_test = np.random.randn(20).astype(np.float32)\n",
    "\n",
    "np.save(os.path.join(train_dir, \"x_train.npy\"), x_train)\n",
    "np.save(os.path.join(train_dir, \"y_train.npy\"), y_train)\n",
    "np.save(os.path.join(test_dir, \"x_test.npy\"), x_test)\n",
    "np.save(os.path.join(test_dir, \"y_test.npy\"), y_test)\n",
    "\n",
    "print(f\"Created training data: {x_train.shape}, {y_train.shape}\")\n",
    "print(f\"Created test data: {x_test.shape}, {y_test.shape}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create pytorch model definition\n",
    "pytorch_model_def = '''\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "def get_model():\n",
    "    return nn.Sequential(\n",
    "        nn.Linear(4, 10),\n",
    "        nn.ReLU(),\n",
    "        nn.Linear(10, 1)\n",
    "    )\n",
    "'''\n",
    "\n",
    "with open(os.path.join(source_dir, \"pytorch_model_def.py\"), 'w') as f:\n",
    "    f.write(pytorch_model_def)\n",
    "\n",
    "print(\"Created pytorch_model_def.py\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create training script (same as single container for simplicity)\n",
    "training_script = '''\n",
    "import argparse\n",
    "import numpy as np\n",
    "import os\n",
    "import sys\n",
    "import logging\n",
    "import json\n",
    "import shutil\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.utils.data import DataLoader, TensorDataset\n",
    "from pytorch_model_def import get_model\n",
    "\n",
    "logger = logging.getLogger(__name__)\n",
    "logger.setLevel(logging.DEBUG)\n",
    "logger.addHandler(logging.StreamHandler(sys.stdout))\n",
    "current_dir = os.path.dirname(os.path.abspath(__file__))\n",
    "data_dir = \"/opt/ml/input/data\"\n",
    "\n",
    "def get_train_data(train_dir):\n",
    "    x_train = np.load(os.path.join(train_dir, \"x_train.npy\"))\n",
    "    y_train = np.load(os.path.join(train_dir, \"y_train.npy\"))\n",
    "    logger.info(f\"x train: {x_train.shape}, y train: {y_train.shape}\")\n",
    "    return torch.from_numpy(x_train), torch.from_numpy(y_train)\n",
    "\n",
    "def get_test_data(test_dir):\n",
    "    x_test = np.load(os.path.join(test_dir, \"x_test.npy\"))\n",
    "    y_test = np.load(os.path.join(test_dir, \"y_test.npy\"))\n",
    "    logger.info(f\"x test: {x_test.shape}, y test: {y_test.shape}\")\n",
    "    return torch.from_numpy(x_test), torch.from_numpy(y_test)\n",
    "\n",
    "def train():\n",
    "    train_dir = os.path.join(data_dir, \"train\")\n",
    "    test_dir = os.path.join(data_dir, \"test\")\n",
    "    model_dir = os.environ.get(\"SM_MODEL_DIR\", os.path.join(current_dir, \"data/model\"))\n",
    "\n",
    "    x_train, y_train = get_train_data(train_dir)\n",
    "    x_test, y_test = get_test_data(test_dir)\n",
    "    train_ds = TensorDataset(x_train, y_train)\n",
    "\n",
    "    batch_size = 64\n",
    "    epochs = 1\n",
    "    learning_rate = 0.1\n",
    "    logger.info(f\"batch_size = {batch_size}, epochs = {epochs}, learning rate = {learning_rate}\")\n",
    "\n",
    "    train_dl = DataLoader(train_ds, batch_size, shuffle=True)\n",
    "    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "    model = get_model().to(device)\n",
    "    criterion = nn.MSELoss()\n",
    "    optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)\n",
    "\n",
    "    for epoch in range(epochs):\n",
    "        for x_train_batch, y_train_batch in train_dl:\n",
    "            y = model(x_train_batch.float())\n",
    "            loss = criterion(y.flatten(), y_train_batch.float())\n",
    "            optimizer.zero_grad()\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "        epoch += 1\n",
    "        logger.info(f\"epoch: {epoch} -> loss: {loss}\")\n",
    "\n",
    "    with torch.no_grad():\n",
    "        y = model(x_test.float()).flatten()\n",
    "        mse = ((y - y_test) ** 2).sum() / y_test.shape[0]\n",
    "    print(\"Test MSE:\", mse.numpy())\n",
    "\n",
    "    os.makedirs(model_dir, exist_ok=True)\n",
    "    torch.save(model.state_dict(), model_dir + \"/model.pth\")\n",
    "    \n",
    "    inference_code_path = model_dir + \"/code/\"\n",
    "    if not os.path.exists(inference_code_path):\n",
    "        os.mkdir(inference_code_path)\n",
    "        logger.info(f\"Created a folder at {inference_code_path}!\")\n",
    "\n",
    "    shutil.copy(\"local_training_script.py\", inference_code_path)\n",
    "    shutil.copy(\"pytorch_model_def.py\", inference_code_path)\n",
    "    logger.info(f\"Saving models files to {inference_code_path}\")\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    print(\"Running the training job ...\")\n",
    "    train()\n",
    "'''\n",
    "\n",
    "with open(os.path.join(source_dir, \"local_training_script.py\"), 'w') as f:\n",
    "    f.write(training_script)\n",
    "\n",
    "print(\"Created local_training_script.py\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 3: Configure Distributed Local Training\n",
    "\n",
    "Set up ModelTrainer for distributed training in local container mode."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "source_code = SourceCode(\n",
    "    source_dir=source_dir,\n",
    "    entry_script=\"local_training_script.py\",\n",
    ")\n",
    "\n",
    "distributed = Torchrun(\n",
    "    process_count_per_node=1,\n",
    ")\n",
    "\n",
    "compute = Compute(\n",
    "    instance_type=\"local_cpu\",\n",
    "    instance_count=2,\n",
    ")\n",
    "\n",
    "train_data = InputData(\n",
    "    channel_name=\"train\",\n",
    "    data_source=train_dir,\n",
    ")\n",
    "\n",
    "test_data = InputData(\n",
    "    channel_name=\"test\",\n",
    "    data_source=test_dir,\n",
    ")\n",
    "\n",
    "print(\"Distributed Local Training Configuration:\")\n",
    "print(f\"  Containers: {compute.instance_count}\")\n",
    "print(f\"  Processes per container: {distributed.process_count_per_node}\")\n",
    "print(f\"  Total processes: {compute.instance_count * distributed.process_count_per_node}\")\n",
    "print(f\"  Distributed framework: Torchrun\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 4: Create Distributed ModelTrainer\n",
    "\n",
    "Initialize ModelTrainer for distributed local container training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_trainer = ModelTrainer(\n",
    "    training_image=DEFAULT_CPU_IMAGE,\n",
    "    sagemaker_session=sagemaker_session,\n",
    "    source_code=source_code,\n",
    "    distributed=distributed,\n",
    "    compute=compute,\n",
    "    input_data_config=[train_data, test_data],\n",
    "    base_job_name=\"local_mode_multi_container\",\n",
    "    training_mode=Mode.LOCAL_CONTAINER,\n",
    ")\n",
    "\n",
    "print(\"Distributed ModelTrainer created successfully!\")\n",
    "print(f\"Training mode: {Mode.LOCAL_CONTAINER}\")\n",
    "print(f\"Distributed config: {distributed}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 5: Run Distributed Local Training\n",
    "\n",
    "Start the distributed training job using multiple local containers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Starting distributed local container training...\")\n",
    "print(\"This will launch 2 Docker containers with 1 training process each.\")\n",
    "\n",
    "try:\n",
    "    model_trainer.train()\n",
    "    print(\"Distributed local container training completed successfully!\")\n",
    "    operation_successful = True\n",
    "except Exception as e:\n",
    "    print(f\"Training failed with error: {e}\")\n",
    "    operation_successful = False"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 6: Check Training Results\n",
    "\n",
    "Examine the results from distributed training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if operation_successful:\n",
    "    current_dir = os.getcwd()\n",
    "    \n",
    "    print(\"Distributed Training Results:\")\n",
    "    print(\"=\" * 35)\n",
    "    \n",
    "    # Check that certain directories don't exist (cleanup verification)\n",
    "    assert not os.path.exists(os.path.join(current_dir, \"shared\"))\n",
    "    assert not os.path.exists(os.path.join(current_dir, \"input\"))\n",
    "    assert not os.path.exists(os.path.join(current_dir, \"algo-1\"))\n",
    "    assert not os.path.exists(os.path.join(current_dir, \"algo-2\"))\n",
    "    print(\"✓ Temporary directories properly cleaned up\")\n",
    "    \n",
    "    # Check for expected artifacts\n",
    "    directories_to_check = [\n",
    "        \"compressed_artifacts\",\n",
    "        \"artifacts\",\n",
    "        \"model\",\n",
    "        \"output\",\n",
    "    ]\n",
    "    \n",
    "    for directory in directories_to_check:\n",
    "        path = os.path.join(current_dir, directory)\n",
    "        if os.path.exists(path):\n",
    "            print(f\"✓ {directory}: Found\")\n",
    "        else:\n",
    "            print(f\"✗ {directory}: Not found\")\n",
    "    \n",
    "    print(\"\\nDistributed Training Configuration:\")\n",
    "    print(f\"  Training Image: {DEFAULT_CPU_IMAGE}\")\n",
    "    print(f\"  Container Count: {compute.instance_count}\")\n",
    "    print(f\"  Processes per Container: {distributed.process_count_per_node}\")\n",
    "    print(f\"  Total Training Processes: {compute.instance_count * distributed.process_count_per_node}\")\n",
    "    \n",
    "else:\n",
    "    print(\"Training was not successful.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 7: Clean Up\n",
    "\n",
    "Clean up local artifacts and temporary files."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    subprocess.run([\"docker\", \"compose\", \"down\", \"-v\"], check=False)\n",
    "    print(\"Docker containers stopped\")\n",
    "except Exception:\n",
    "    pass\n",
    "\n",
    "# Clean up temporary files\n",
    "try:\n",
    "    shutil.rmtree(temp_dir)\n",
    "    print(f\"Cleaned up temporary directory: {temp_dir}\")\n",
    "except Exception as e:\n",
    "    print(f\"Could not clean up temp directory: {e}\")\n",
    "\n",
    "# Clean up training artifacts\n",
    "current_dir = os.getcwd()\n",
    "directories = [\"compressed_artifacts\", \"artifacts\", \"model\", \"output\"]\n",
    "\n",
    "for directory in directories:\n",
    "    path = os.path.join(current_dir, directory)\n",
    "    if os.path.exists(path):\n",
    "        try:\n",
    "            shutil.rmtree(path)\n",
    "            print(f\"Cleaned up: {directory}\")\n",
    "        except Exception as e:\n",
    "            print(f\"Could not clean up {directory}: {e}\")\n",
    "\n",
    "# Final assertion\n",
    "assert operation_successful\n",
    "print(\"\\n✓ Distributed local training completed successfully!\")\n",
    "print(\"Cleanup completed - all artifacts removed.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "This notebook demonstrated:\n",
    "1. **Multi-container distributed training**: Running training across multiple Docker containers locally\n",
    "2. **Torchrun integration**: Using SageMaker's Torchrun distributed driver\n",
    "3. **Local development workflow**: Testing distributed training before cloud deployment\n",
    "4. **Proper cleanup**: Following cleanup patterns for local artifacts\n",
    "\n",
    "Distributed local training provides a great way to test distributed training patterns locally before deploying to SageMaker cloud instances, with no AWS costs and realistic container-based execution."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv-test",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
